!-----------------------------------------------------------------------
! This is a 1-D Monte Carlo solver that transportts neutrons across a
! slab. Vacuum boundry conditions on each side of the slab. This problem
! tallies the absorption, transmission and reflection probabilities.
! Aaron Bevill
! Mateusz Monterial
! 10/20/2013
!-----------------------------------------------------------------------


program MCsolve
IMPLICIT NONE
! neutron properties
double precision :: x, E, mu                ! position, energy, directional cosine
double precision :: xn, En, mu_n            ! new pos, energy, directional cosine
double precision :: r, theta, phi           ! distance traveled, scattering, polar angles
! problem/material parameters
double precision :: A, alph, beta, D=1      ! mass number, thickness of slab
integer :: nmat                             ! number of materials
integer :: imodel                           ! model iterator
double precision, dimension(:), allocatable :: micro_t, micro_s ! nuclear interaction data per atom
double precision :: Sig_t, Sig_s            ! nuclear interaction data
double precision, dimension(:), allocatable :: unper_dens ! unperturbed atom densities
double precision, dimension(:,:), allocatable :: per_dens ! perturbed atom densities
! random number properties
integer*8 :: irand=0                    ! random number counter
integer*8 :: max_stride=0, stride=150   ! N random numbers reserved per history
double precision :: rand1, rand2        ! random num storage
integer, parameter :: seed = 86456      ! seed for random num gen
! trace data log
integer, parameter :: segsize=4000 ! allocated number of segments per batch
integer, parameter :: nbatch=100 ! number of hist to run before retracing
double precision, dimension(:,:), allocatable :: log_micro_t, log_micro_s
double precision, dimension(segsize) :: log_sig_t, log_sig_s, log_dtc
integer, dimension(segsize) :: log_loss
integer :: iseg, nseg, max_seg=0! segment iterator and max tracker
! other variable
double precision :: ref=0, trn=0, abs=0     ! reflected and transition tallies
integer :: iloss                ! loss type (reflec'n/transmission/absorption)
double precision, dimension(:,:), allocatable :: retrace_tallies ! tallies in the perturbed models
integer*8 :: NPS, i                     ! number of particles simulated
character(32) :: NPS_buff, nmat_buff        ! buffers for reading in nps and nmat
double precision :: wgt                     ! weight of the retraced history
double precision :: sensitivity_delta=0.1   ! the perturbation quantity
logical, parameter :: retrace_flag=.true.  ! retrace histories iff true
double precision :: t_begin, t_end        ! timers

! get the problem size from input line
call getarg(1, nmat_buff) ! get the number of particles
call getarg(2, NPS_buff)  ! get the number of iterations
! convert buffers to integers
read(nmat_buff, '(i10)') nmat
if (nmat == 0) then
    nmat = 100
end if
read(NPS_buff, '(i10)') NPS
if (NPS == 0) then
    NPS = 1e5
end if
! allocate various arrays
allocate(micro_t(nmat))
allocate(micro_s(nmat))
allocate(unper_dens(nmat))
allocate(per_dens(nmat,nmat))
allocate(log_micro_t(nmat,segsize))
allocate(log_micro_s(nmat,segsize))
allocate(retrace_tallies(nmat,3))

! load material properties (no energy dependence)
! add the read in ability so these are not changed
A = 200.
alph = ((A-1.)/(A+1.))**2. ! calculare alpha
beta = 1. - alph

! fill the atom_densities matrix with unperturbed and perturbed densities
call init_atom_densities(nmat, sensitivity_delta, unper_dens, per_dens)

! main do loop to simulate the particles transport through slab
call cpu_time(t_begin)
call srand(seed) ! initilize the random number generator
iseg = 0
do i = 0, NPS-1
    ! initilize neutron properties
    x=0.0           ! position
    E=1.0           ! energy
    mu=1.0          ! incident angle to slab
    call advance_rng(i, irand, stride, max_stride)
    iseg = iseg + 1
    segment_loop: do iseg=iseg,segsize
        ! sample random nums for distance and absoprtion
        rand1 = rand()
        rand2 = rand()
        irand = irand + 2
        ! calculate new parameters
        call get_micro_xs(E, nmat, micro_t, micro_s)
        Sig_t = dot_product(micro_t, unper_dens)
        Sig_s = dot_product(micro_s, unper_dens)
        r = -1./Sig_t*dlog(rand1)       ! sample distance traveled
        xn = x + r*mu                   ! calculate new position
        ! log the information for later retracing
        if (retrace_flag) then
            log_micro_t(:, iseg) = micro_t
            log_micro_s(:, iseg) = micro_s
            log_sig_t(iseg) = Sig_t
            log_sig_s(iseg) = Sig_s
            log_dtc(iseg) = r
        end if
        ! check the boundry conditions
        if (xn<0.) then                 ! the particle was reflected
            ref = ref+1.
            if (retrace_flag) then
                log_loss(iseg) = 1
                log_dtc(iseg) = (0. - x) / mu ! distance to escape
            end if
            exit segment_loop
        else if (xn>D) then             ! the particle was transmited
            trn = trn+1.
            if (retrace_flag) then
                log_loss(iseg) = 2
                log_dtc(iseg) = (D - x) / mu ! distance to escape
            end if
            exit segment_loop
        ! check if absorbed
        else if (rand2 * Sig_t > Sig_s) then
            abs = abs+1.
            if (retrace_flag) then
                log_loss(iseg) = 3
            end if
            exit segment_loop
        ! neutrons scatters inside the slab
        else
            rand1 = rand()
            rand2 = rand()
            irand = irand + 2
            ! calculate new energy
            En = E - rand1 * beta * E
            ! calculate new direction cosine
            mu_n = 2 * rand2 - 1
            ! save the new parameters
            E = En
            mu = mu_n
            x = xn
            if (retrace_flag) then
                log_loss(iseg) = 0
            end if
        end if
    end do segment_loop
    ! retrace nbatch histories in a batch
    if (retrace_flag .and. mod(i, nbatch) == nbatch-1) then
        nseg = iseg ! N segments in this batch
        max_seg = max(max_seg, nseg)
        wgt = 1
        do imodel = 1, nmat
            do iseg = 1, nseg
                Sig_t = dot_product(log_micro_t(:, iseg), per_dens(:, imodel))
                Sig_s = dot_product(log_micro_s(:, iseg), per_dens(:, imodel))
                iloss = log_loss(iseg)
                if (iloss == 0) then
                    wgt = wgt * Sig_s / log_sig_s(iseg) &
                        * dexp(log_dtc(iseg) * (log_sig_t(iseg) - Sig_t))
                ! If escaped, correct by prob of escape
                else if (iloss < 3) then
                    wgt = wgt * dexp(log_dtc(iseg) * (log_sig_t(iseg) - Sig_t))
                    retrace_tallies(imodel, iloss) = &
                        retrace_tallies(imodel, iloss) + wgt
                    wgt = 1
                ! If absorbed, use that correction
                else if (iloss == 3) then
                    wgt = wgt * (Sig_t-Sig_s) / (log_sig_t(iseg)-log_sig_s(iseg)) &
                        * dexp(log_dtc(iseg) * (log_sig_t(iseg) - Sig_t))
                    retrace_tallies(imodel, iloss) = &
                        retrace_tallies(imodel, iloss) + wgt
                    wgt = 1
                end if
            end do
        end do
        iseg = 0
    end if
end do
call cpu_time(t_end)
! print the result of the tallies
write(*,*) 'P(ref)', ref/real(nps)
write(*,*) 'P(trn)', trn/real(nps)
write(*,*) 'P(abs)', abs/real(nps)
write(*,*) 'max segments/batch', max_seg
write(*,*) 'max stride', max_stride
write(*,*) 'Time (seconds)', t_end - t_begin
if (retrace_flag) then
    write(*,*) 'Sensitivity checksums'
    write(*,"(4X,ES22.16)") sum(retrace_tallies(:,1))
    write(*,"(4X,ES22.16)") sum(retrace_tallies(:,2))
    write(*,"(4X,ES22.16)") sum(retrace_tallies(:,3))
    if (nmat == 100 .and. NPS == 1e5) then
        call internal_check(nmat, retrace_tallies, 1e-15)
    end if
    !write(*,*) "Sensitivity estimates"
    !do imodel = 1, nmat
    !    write(*,"(4X,4ES15.7)") &
    !        (retrace_tallies(imodel, 1) - ref) / real(NPS) / sensitivity_delta, &
    !        (retrace_tallies(imodel, 2) - trn) / real(NPS) / sensitivity_delta, &
    !        (retrace_tallies(imodel, 3) - abs) / real(NPS) / sensitivity_delta ! , &
    !       ! (sum(retrace_tallies(imodel, :)) - ref - trn - abs) &
    !        !    / real(NPS) / sensitivity_delta
    !end do
end if
end program MCsolve


subroutine init_atom_densities(nmat, sensitivity_delta, unper_dens, per_dens)
    implicit none
    integer, intent(in) :: nmat
    real, intent(in) :: sensitivity_delta
    real, intent(out), dimension(nmat) :: unper_dens
    real, intent(out), dimension(nmat,nmat) :: per_dens
    !
    integer :: imat, imodel, linspace(nmat)

    linspace = (/(imat, imat=1,nmat)/)
    unper_dens = 0.02 * dexp(-0.001 * linspace)
    do imodel=1,nmat
        per_dens(:, imodel) = unper_dens(:)
        per_dens(imodel, imodel) = per_dens(imodel, imodel) * &
            (1. + sensitivity_delta)
    end do
end subroutine

subroutine get_micro_xs(energy, nmat, micro_t, micro_s)
    implicit none
    real, intent(in)  :: energy
    integer, intent(in) :: nmat
    real, dimension(nmat), intent(out) :: micro_t, micro_s
    !
    integer :: linspace(nmat), imat

    linspace = (/(imat, imat=1,nmat)/)
    micro_t = dsin(linspace * dlog(energy)) + 1.1
    micro_s = micro_t * 0.8 * (dsin(4*linspace*dlog(energy)) + 1)
end subroutine get_micro_xs

subroutine advance_rng(ihist, irand, stride, max_stride)
    implicit none
    integer*8, intent(in) :: ihist, stride
    integer*8, intent(inout) :: irand, max_stride
    !
    double precision :: tmp

    if (irand > 0) then
        max_stride = max(max_stride, irand - stride * (ihist - 1))
    end if
    if (irand > stride * ihist) then
        write(*,*) "WARNING: Stride exceeded."
    end if
    do irand = irand, stride * ihist - 1
        tmp = rand()
    end do
end subroutine

subroutine internal_check(nmat, retrace_tallies, eps)
    implicit none
    integer, intent(in) :: nmat
    double precision, dimension(nmat, 3), intent(in) :: retrace_tallies
    double precision, intent(in) :: eps
    !
    double precision :: s
    logical :: wrong_flag=.false.
    character(LEN=*), PARAMETER :: CONSOLE_GREEN=achar(27)//"[1;32m"
    character(LEN=*), PARAMETER :: CONSOLE_RED=achar(27)//"[1;31m"
    character(LEN=*), PARAMETER :: CONSOLE_DEFAULT=achar(27)//"[0m"

    s = sum(retrace_tallies(:,1))
    wrong_flag = wrong_flag .or. dabs(s / 2.2648737035819185E+06 - 1) > eps
    s = sum(retrace_tallies(:,2))
    wrong_flag = wrong_flag .or. dabs(s / 2.5622963911091546E+06 - 1) > eps
    s = sum(retrace_tallies(:,3))
    wrong_flag = wrong_flag .or. dabs(s / 5.1728524632014949E+06 - 1) > eps
    !
    if (wrong_flag) then
        write(*,*) CONSOLE_RED, "INTERNAL CHECK FAIL", CONSOLE_DEFAULT
    else
        write(*,*) CONSOLE_GREEN, "INTERNAL CHECK PASS", CONSOLE_DEFAULT
    end if
end subroutine internal_check
